# -*- coding: UTF-8 -*-
# Author: Liang Bing
# E-mail: believer19940901@gmail.com
# Date Time: 2025/10/28 13:00
"""
脚本功能：
---------
将 RefSeq 的 GRCh37 基因组序列的染色体名称修改为 UCSC 格式（chr 前缀），
并去除非标准染色体（如 unplaced、alternate、random 等）。
"""

from Bio import SeqIO
from argparse import ArgumentParser, Namespace
from pathlib import Path
import re


def _parse_chr_name(description: str) -> str | None:
    """
    从 RefSeq 描述行中解析染色体名称并转换为 UCSC 格式。

    例如：
    >gi|568336023|ref|NC_000001.11| Homo sapiens chromosome 1, GRCh37.p13 Primary Assembly
    -> chr1

    参数:
        description (str): 序列描述信息。

    返回:
        str | None: 转换后的染色体名称（如 chr1），未匹配则返回 None。
    """
    match = re.search(
        r"Homo\s+sapiens\s+chromosome\s+(\S+),\s+\S+\s+Primary\s+Assembly",
        description
    )
    if not match:
        print(f"[INFO] Skipped: {description}")
        return None

    chrom = match.group(1)
    chrom = chrom.replace("MT", "M")  # 兼容线粒体命名
    new_name = f"chr{chrom}"
    print(f"[INFO] Parsed: {description} -> {new_name}")
    return new_name


def main(args: Namespace) -> None:
    """
    主函数：读取输入 fasta 文件，修改染色体名称并输出。
    """
    input_path = Path(args.input)
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    print(f"[START] Reading {input_path}")
    print(f"[INFO] Writing modified sequences to {output_path}")

    out_records = []
    with open(input_path, "r", encoding="utf-8") as fr:
        for record in SeqIO.parse(fr, "fasta"):
            new_name = _parse_chr_name(record.description)
            if new_name:
                record.id = new_name
                record.name = new_name
                record.description = new_name
                out_records.append(record)

    SeqIO.write(out_records, output_path, "fasta")
    print(f"[DONE] Total {len(out_records)} sequences written to {output_path}")


if __name__ == "__main__":
    parser = ArgumentParser(description="Convert RefSeq GRCh37 chromosome names to UCSC-style (chr prefix)")
    parser.add_argument("-i", "--input", type=str, required=True, help="Input FASTA file (RefSeq GRCh37)")
    parser.add_argument("-o", "--output", type=str, required=True, help="Output FASTA file with UCSC chromosome names")
    main(parser.parse_args())
